package com.aocore.proxy;

import com.aocore.awesomo.AwesomO;
import com.aocore.proxy.data.ConnectionData;
import com.aocore.packet.Packet;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import org.netbeans.lib.collab.util.Socks5SocketChannelAdaptor;

/**
 *
 * @author derek
 */
public abstract class Proxy extends Thread
{
    public static Map<Integer, AwesomO> bots = new HashMap<Integer, AwesomO>();

    protected AwesomO awesomO = null;
    private Selector selector = null;
    private SocketChannel serverSocketChannel;
    private SocketChannel clientSocketChannel;

    private ConnectionData serverConnectionData = new ConnectionData()
    {
        @Override
        public void callback() throws IOException
        {
            for (Packet packet : parseClientPacket(read))
            {
                Packet p = processClientPacket(packet);

                if (p != null && awesomO != null)
                {
                    p = awesomO.processPacket(p);
                }

                if (p != null)
                {
                    writeClientPacket(p);
                }
            }
        }
    };

    private ConnectionData clientConnectionData = new ConnectionData()
    {
        @Override
        public void callback() throws IOException
        {
            for (Packet packet : parseServerPacket(read))
            {
                Packet p = processServerPacket(packet);

                if (p != null && awesomO != null)
                {
                    p = awesomO.processPacket(p);
                }

                if (p != null)
                {
                    writeServerPacket(p);
                }
            }
        }
    };

    public Proxy(SocketChannel socketChannel) throws IOException
    {
        // create a new selector
        selector = Selector.open();

        // register non blocking server socket
        serverSocketChannel = socketChannel;
        serverSocketChannel.configureBlocking(false);
        serverSocketChannel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE, serverConnectionData);
    }

    protected abstract List<Packet> parseClientPacket(ByteBuffer byteBuffer);
    protected abstract List<Packet> parseServerPacket(ByteBuffer byteBuffer);
    protected abstract Packet processClientPacket(Packet packet);
    protected abstract Packet processServerPacket(Packet packet);

    public void connectClient(SocketAddress remote) throws IOException
    {
        if (false)
        {
            clientSocketChannel = Socks5SocketChannelAdaptor.open(remote, new InetSocketAddress("71.21.135.166", 1469));
            clientSocketChannel.configureBlocking(false);
            clientSocketChannel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE, clientConnectionData);
        }
        else
        {
            // register non blocking client socket
            clientSocketChannel = SocketChannel.open();
            clientSocketChannel.configureBlocking(false);
            clientSocketChannel.connect(remote);
            clientSocketChannel.register(selector, SelectionKey.OP_CONNECT, clientConnectionData);
        }
    }

    protected void writeClient(byte... data) throws IOException
    {
        clientConnectionData.write(data);

        if (clientSocketChannel != null && clientSocketChannel.isConnected())
        {
            clientSocketChannel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE, clientConnectionData);
            selector.wakeup();
        }
    }

    protected void writeClientPacket(Packet packet) throws IOException
    {
        clientConnectionData.writePacket(packet);

        if (clientSocketChannel != null && clientSocketChannel.isConnected())
        {
            clientSocketChannel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE, clientConnectionData);
            selector.wakeup();
        }
    }

    protected void writeServer(byte... data) throws IOException
    {
        serverConnectionData.write(data);

        if (serverSocketChannel != null && serverSocketChannel.isConnected())
        {
            serverSocketChannel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE, serverConnectionData);
            selector.wakeup();
        }
    }

    protected void writeServerPacket(Packet packet) throws IOException
    {
        serverConnectionData.writePacket(packet);

        if (serverSocketChannel != null && serverSocketChannel.isConnected())
        {
            serverSocketChannel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE, serverConnectionData);
            selector.wakeup();
        }
    }

    public void stopProxy()
    {
        System.out.println("STOP PROXY");

        try
        {
            if (serverSocketChannel != null)
            {
                serverSocketChannel.close();
            }
        }
        catch (Exception e)
        {
            e.printStackTrace(System.err);
        }

        try
        {
            if (clientSocketChannel != null)
            {
                clientSocketChannel.close();
            }
        }
        catch (Exception e)
        {
            e.printStackTrace(System.err);
        }

        try
        {
            if (selector.isOpen())
            {
                selector.close();
            }
        }
        catch (Exception e)
        {
            e.printStackTrace(System.err);
        }
/*
        if (awesomO != null)
        {
            if (this instanceof ChatProxy)
            {
                awesomO.shutdown();
            }

            awesomO.interrupt();
        }*/
    }

    @Override
    public final void run()
    {
        try
        {
            while (selector.select() >= 0)
            {
                Iterator<SelectionKey> it = selector.selectedKeys().iterator();

                while (it.hasNext())
                {
                    SelectionKey key = it.next();
                    it.remove();

                    if (key.isConnectable())
                    {
                        //System.out.println("CLIENT CONNECTED");
                        SocketChannel socketChannel = (SocketChannel) key.channel();
                        socketChannel.finishConnect();
                        socketChannel.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE, key.attachment());
                    }

                    if (key.isWritable())
                    {
                        //System.out.println("WRITEABLE");
                        ConnectionData connection = (ConnectionData) key.attachment();
                        int length = connection.write((SocketChannel) key.channel());

                        if (length == 0)
                        {
                            SocketChannel socketChannel = (SocketChannel) key.channel();
                            socketChannel.register(selector, SelectionKey.OP_READ, key.attachment());
                        }
                    }

                    if (key.isReadable())
                    {
                        //System.out.println("READABLE");
                        ConnectionData connection = (ConnectionData) key.attachment();
                        int length = connection.read((SocketChannel) key.channel());

                        if (length > 0)
                        {
                            connection.callback();
                        }
                        else if (length == -1)
                        {
                            stopProxy();
                            return;
                        }
                    }
                }
            }
        }
        catch (Exception e)
        {
            e.printStackTrace(System.err);
            stopProxy();
            return;
        }
    }
}
